package com.example.scheduler;

import com.example.mq.RocketMqProducer;
import com.example.mq.RocketMqConsumerService;
import com.example.raft.RaftNode;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.common.message.MessageExt;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * SchedulerNode：在 RaftNode 基础上扩展调度能力
 *  - slowTaskMap：统计普通任务耗时，若>500ms多次则标记慢任务
 *  - dagTaskMap：管理正在执行的DAG任务
 *  - scheduleDag()：可使用DagTemplate进行DAG拆分并逐层调度
 */
public class SchedulerNode extends RaftNode {

    private final RocketMqProducer mqProducer;
    private final RocketMqConsumerService consumerService;
    private final List<TaskExecutor> executors;

    // 是否处于高峰期
    private volatile boolean peakMode = false;

    // 普通任务慢任务统计表 jobId->超时次数
    private final ConcurrentHashMap<String, AtomicInteger> slowTaskMap = new ConcurrentHashMap<>();
    // DAG管理表 dagId->DagTask
    private final ConcurrentHashMap<String, DagTask> dagTaskMap = new ConcurrentHashMap<>();

    // 用于监听执行器发送到 RESULT_TOPIC 的结果消息
    private DefaultMQPushConsumer resultConsumer;

    public SchedulerNode(String nodeId,
                         List<String> peerIds,
                         RocketMqProducer mqProducer,
                         RocketMqConsumerService consumerService,
                         List<TaskExecutor> executors) {
        super(nodeId, peerIds);
        this.mqProducer = mqProducer;
        this.consumerService = consumerService;
        this.executors = executors;
    }

    /**
     * 启动调度器的MQ消费者，用于接收执行器的任务执行结果（RESULT_TOPIC）
     */
    public void startScheduler() {
        try {
            resultConsumer = consumerService.startConsumer(
                    "scheduler-result-group-" + nodeId,
                    "RESULT_TOPIC",
                    (List<MessageExt> msgs, ConsumeConcurrentlyContext context) -> {
                        handleResultMessages(msgs);
                        return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
                    }
            );
            System.out.println("Scheduler " + nodeId + " startScheduler: MQ consumer started.");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 处理执行器发回的“RESULT|OK|jobId|COST=xxx”消息
     * 若jobId含冒号，则视为DAG子任务；否则视为普通任务
     */
    private void handleResultMessages(List<MessageExt> msgs) {
        if (getState() != NodeState.LEADER) {
            return;
        }
        for (MessageExt msg : msgs) {
            try {
                String body = new String(msg.getBody(), "UTF-8");
                if (body.startsWith("RESULT|")) {
                    String[] parts = body.split("\\|");
                    if (parts.length >= 3) {
                        String resultStatus = parts[1]; // OK or FAIL
                        String jobId = parts[2];
                        long cost = 0;
                        if (parts.length >= 4 && parts[3].startsWith("COST=")) {
                            cost = Long.parseLong(parts[3].substring("COST=".length()));
                        }
                        if (jobId.contains(":")) {
                            // DAG子任务结果
                            processDagResult(jobId, cost, resultStatus);
                        } else {
                            // 普通任务，更新慢任务统计
                            System.out.println("Scheduler " + nodeId + " received result => job=" + jobId +
                                    ", cost=" + cost + "ms, status=" + resultStatus);
                            if ("OK".equals(resultStatus) && cost > 500) {
                                slowTaskMap.computeIfAbsent(jobId, k -> new AtomicInteger(0))
                                        .incrementAndGet();
                            }
                        }
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 设置是否高峰期（peakMode）
     */
    public void setPeakMode(boolean pm) {
        this.peakMode = pm;
        System.out.println("Scheduler " + nodeId + " set peakMode=" + pm);
    }

    /**
     * 调度一个普通任务
     * 根据 slowTaskMap 判断其是否慢任务 (isKnownSlow)，再决定直连执行器或投MQ
     */
    public void scheduleTask(String jobId) {
        if (getState() != NodeState.LEADER) {
            System.out.println("Scheduler " + nodeId + " is not leader, skip scheduling: " + jobId);
            return;
        }
        AtomicInteger countObj = slowTaskMap.get(jobId);
        boolean isKnownSlow = (countObj != null && countObj.get() >= 10);

        if (!peakMode) {
            // 普通期，直接RPC调度执行器
            if (executors.isEmpty()) {
                System.out.println("No executors available for job: " + jobId);
            } else {
                TaskExecutor executor = executors.get(0);
                System.out.println("Scheduler " + nodeId + " => direct call to " +
                        executor.getExecutorId() + ", job=" + jobId + ", isSlow=" + isKnownSlow);
                executor.directExecute(jobId, isKnownSlow);
            }
        } else {
            // 高峰期，将任务投递到 MQ
            String prefix = isKnownSlow ? "TASK|SLOW|" : "TASK|FAST|";
            try {
                System.out.println("Scheduler " + nodeId + " => dispatch to MQ: " + prefix + jobId);
                mqProducer.sendMessage("TASK_TOPIC", prefix + jobId);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 调度一个DAG任务。可从DagTemplateManager中获取对应模板
     */
    public void scheduleDag(String dagId, DagTemplate template) {
        if (getState() != NodeState.LEADER) {
            System.out.println("Scheduler " + nodeId + " is not leader, cannot schedule DAG: " + dagId);
            return;
        }
        try {
            // 通过模板进行拓扑排序，分层
            List<List<String>> levels = template.getLevels();
            DagTask dagTask = new DagTask(dagId, levels);
            dagTaskMap.put(dagId, dagTask);

            System.out.println("Scheduler " + nodeId + " scheduling DAG " + dagId
                    + " with levels: " + levels);

            // 调度第一层
            List<String> firstLevel = dagTask.getCurrentLevelTasks();
            if (firstLevel != null) {
                for (String subTaskId : firstLevel) {
                    scheduleDagSubTask(dagId, subTaskId, false);
                }
            }
        } catch (Exception e) {
            System.out.println("Error scheduling DAG " + dagId + ": " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 调度DAG子任务（"dagId:subTaskId"）
     */
    private void scheduleDagSubTask(String dagId, String subTaskId, boolean isSlow) {
        String fullTaskId = dagId + ":" + subTaskId;
        if (!peakMode) {
            if (executors.isEmpty()) {
                System.out.println("No executors available for DAG sub-task: " + fullTaskId);
            } else {
                TaskExecutor executor = executors.get(0);
                System.out.println("Scheduler " + nodeId + " => direct call executor for DAG subTask: "
                        + fullTaskId + ", isSlow=" + isSlow);
                executor.directExecute(fullTaskId, isSlow);
            }
        } else {
            String prefix = isSlow ? "TASK|SLOW|" : "TASK|FAST|";
            try {
                System.out.println("Scheduler " + nodeId + " => dispatch DAG subTask to MQ: " + fullTaskId);
                mqProducer.sendMessage("TASK_TOPIC", prefix + fullTaskId);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 处理DAG子任务执行结果
     */
    public void processDagResult(String jobId, long cost, String result) {
        if (!jobId.contains(":")) {
            // 不是DAG任务
            return;
        }
        String[] parts = jobId.split(":");
        String dagId = parts[0];
        DagTask dagTask = dagTaskMap.get(dagId);
        if (dagTask == null) {
            System.out.println("Scheduler " + nodeId + " => DAG " + dagId + " not found?");
            return;
        }
        // 标记该子任务完成
        dagTask.markTaskFinished(jobId);
        System.out.println("Scheduler " + nodeId + " updated DAG " + dagId
                + " => sub-task " + jobId + " done, cost=" + cost + "ms, result=" + result);

        // 若当前层已全部完成，启动下一层
        if (dagTask.isCurrentLevelComplete()) {
            System.out.println("DAG " + dagId + " current level complete.");
            if (dagTask.moveToNextLevel()) {
                List<String> nextLevel = dagTask.getCurrentLevelTasks();
                System.out.println("Scheduler " + nodeId + " scheduling next-level tasks for DAG " + dagId
                        + ": " + nextLevel);
                for (String subTaskId : nextLevel) {
                    scheduleDagSubTask(dagId, subTaskId, false);
                }
            } else {
                System.out.println("DAG " + dagId + " fully completed!");
                dagTaskMap.remove(dagId);
            }
        }
    }

    @Override
    protected void persistLogEntryToDB(LogEntry entry) {
        super.persistLogEntryToDB(entry);
        System.out.println("Scheduler " + nodeId + " => [Redis Del] DelayedDoubleDelete key=" + entry.getCommand());
    }

    /**
     * 停止调度器（关闭MQ消费者）
     */
    public void stopScheduler() {
        if (resultConsumer != null) {
            resultConsumer.shutdown();
        }
        System.out.println("Scheduler " + nodeId + " stopScheduler called.");
    }
}